{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "a23a42c1-66ff-48ce-9bc1-c65bd23a361a",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Package                   Version\n",
      "------------------------- --------------------\n",
      "absl-py                   2.1.0\n",
      "accelerate                1.4.0\n",
      "aiohappyeyeballs          2.5.0\n",
      "aiohttp                   3.11.13\n",
      "aiosignal                 1.3.2\n",
      "albucore                  0.0.23\n",
      "albumentations            1.4.24\n",
      "annotated-types           0.7.0\n",
      "antlr4-python3-runtime    4.9.3\n",
      "anyio                     4.6.2\n",
      "argon2-cffi               23.1.0\n",
      "argon2-cffi-bindings      21.2.0\n",
      "arrow                     1.3.0\n",
      "astor                     0.8.1\n",
      "asttokens                 3.0.0\n",
      "async-lru                 2.0.5\n",
      "async-timeout             5.0.1\n",
      "attrdict                  2.0.1\n",
      "attrs                     24.3.0\n",
      "babel                     2.16.0\n",
      "bce-python-sdk            0.9.29\n",
      "beautifulsoup4            4.12.3\n",
      "black                     25.1.0\n",
      "bleach                    6.2.0\n",
      "blinker                   1.9.0\n",
      "boto3                     1.37.9\n",
      "botocore                  1.37.9\n",
      "braceexpand               0.1.7\n",
      "Brotli                    1.0.9\n",
      "cachetools                5.5.2\n",
      "certifi                   2025.1.31\n",
      "cffi                      1.17.1\n",
      "charset-normalizer        3.3.2\n",
      "click                     8.1.8\n",
      "cloudpickle               3.1.1\n",
      "colorama                  0.4.6\n",
      "coloredlogs               15.0.1\n",
      "colorlog                  6.9.0\n",
      "comm                      0.2.2\n",
      "contourpy                 1.3.1\n",
      "cryptography              44.0.2\n",
      "cssselect                 1.3.0\n",
      "cssutils                  2.11.1\n",
      "cycler                    0.12.1\n",
      "Cython                    3.0.12\n",
      "datasets                  3.3.2\n",
      "debugpy                   1.8.13\n",
      "decorator                 5.1.1\n",
      "defusedxml                0.7.1\n",
      "detectron2                0.6\n",
      "dill                      0.3.8\n",
      "distro                    1.9.0\n",
      "doclayout_yolo            0.0.2b1\n",
      "einops                    0.8.1\n",
      "et_xmlfile                2.0.0\n",
      "evaluate                  0.4.3\n",
      "exceptiongroup            1.2.0\n",
      "executing                 2.2.0\n",
      "fairscale                 0.4.13\n",
      "fast-langdetect           0.2.5\n",
      "fastjsonschema            2.21.1\n",
      "fasttext-predict          0.9.2.4\n",
      "filelock                  3.13.1\n",
      "fire                      0.7.0\n",
      "Flask                     3.1.0\n",
      "flask-babel               4.0.0\n",
      "flatbuffers               25.2.10\n",
      "fonttools                 4.56.0\n",
      "fqdn                      1.5.1\n",
      "frozenlist                1.5.0\n",
      "fsspec                    2024.6.1\n",
      "ftfy                      6.3.1\n",
      "future                    1.0.0\n",
      "fvcore                    0.1.5.post20221221\n",
      "grpcio                    1.70.0\n",
      "h11                       0.14.0\n",
      "httpcore                  1.0.2\n",
      "httpx                     0.27.0\n",
      "huggingface-hub           0.29.2\n",
      "humanfriendly             10.0\n",
      "hydra-core                1.3.2\n",
      "idna                      3.10\n",
      "imageio                   2.37.0\n",
      "imgaug                    0.4.0\n",
      "intel-openmp              2021.4.0\n",
      "iopath                    0.1.9\n",
      "ipykernel                 6.29.5\n",
      "ipython                   8.34.0\n",
      "ipywidgets                8.1.5\n",
      "isoduration               20.11.0\n",
      "itsdangerous              2.2.0\n",
      "jedi                      0.19.2\n",
      "Jinja2                    3.1.4\n",
      "jiter                     0.8.2\n",
      "jmespath                  1.0.1\n",
      "joblib                    1.4.2\n",
      "json5                     0.10.0\n",
      "jsonpointer               3.0.0\n",
      "jsonschema                4.23.0\n",
      "jsonschema-specifications 2024.10.1\n",
      "jupyter                   1.1.1\n",
      "jupyter_client            8.6.3\n",
      "jupyter-console           6.6.3\n",
      "jupyter_core              5.7.2\n",
      "jupyter-events            0.12.0\n",
      "jupyter-lsp               2.2.5\n",
      "jupyter_server            2.15.0\n",
      "jupyter_server_terminals  0.5.3\n",
      "jupyterlab                4.3.6\n",
      "jupyterlab_pygments       0.3.0\n",
      "jupyterlab_server         2.27.3\n",
      "jupyterlab_widgets        3.0.13\n",
      "kiwisolver                1.4.8\n",
      "lazy_loader               0.4\n",
      "lmdb                      1.6.2\n",
      "loguru                    0.7.3\n",
      "lxml                      5.3.1\n",
      "magic-pdf                 1.2.2\n",
      "Markdown                  3.7\n",
      "MarkupSafe                2.1.5\n",
      "matplotlib                3.9.0\n",
      "matplotlib-inline         0.1.7\n",
      "mistune                   3.1.2\n",
      "mkl                       2021.4.0\n",
      "more-itertools            10.6.0\n",
      "mpmath                    1.3.0\n",
      "multidict                 6.1.0\n",
      "multiprocess              0.70.16\n",
      "mypy-extensions           1.0.0\n",
      "nbclient                  0.10.2\n",
      "nbconvert                 7.16.6\n",
      "nbformat                  5.10.4\n",
      "nest-asyncio              1.6.0\n",
      "networkx                  3.3\n",
      "notebook                  7.3.3\n",
      "notebook_shim             0.2.4\n",
      "numpy                     1.26.3\n",
      "omegaconf                 2.3.0\n",
      "onnxruntime               1.21.0\n",
      "openai                    1.65.5\n",
      "opencv-contrib-python     4.6.0.66\n",
      "opencv-python             4.6.0.66\n",
      "opencv-python-headless    4.11.0.86\n",
      "openpyxl                  3.1.5\n",
      "opt-einsum                3.3.0\n",
      "overrides                 7.7.0\n",
      "packaging                 24.2\n",
      "paddleocr                 2.7.3\n",
      "paddlepaddle              2.6.1\n",
      "paddlepaddle-gpu          2.6.1\n",
      "pandas                    2.2.3\n",
      "pandocfilters             1.5.1\n",
      "parso                     0.8.4\n",
      "pathspec                  0.12.1\n",
      "pdf2docx                  0.5.8\n",
      "pdfminer.six              20231228\n",
      "pillow                    11.0.0\n",
      "pip                       25.0\n",
      "platformdirs              3.10.0\n",
      "portalocker               3.1.1\n",
      "premailer                 3.10.0\n",
      "prometheus_client         0.21.1\n",
      "prompt_toolkit            3.0.50\n",
      "propcache                 0.3.0\n",
      "protobuf                  3.20.2\n",
      "psutil                    5.9.0\n",
      "pure_eval                 0.2.3\n",
      "py-cpuinfo                9.0.0\n",
      "pyarrow                   19.0.1\n",
      "pyclipper                 1.3.0.post6\n",
      "pycocotools               2.0.8\n",
      "pycparser                 2.21\n",
      "pycryptodome              3.21.0\n",
      "pydantic                  2.10.6\n",
      "pydantic_core             2.27.2\n",
      "Pygments                  2.19.1\n",
      "PyMuPDF                   1.24.14\n",
      "pyparsing                 3.2.1\n",
      "PyQt6                     6.7.1\n",
      "PyQt6_sip                 13.9.1\n",
      "pyreadline3               3.5.4\n",
      "PySocks                   1.7.1\n",
      "python-dateutil           2.9.0.post0\n",
      "python-docx               1.1.2\n",
      "python-json-logger        3.3.0\n",
      "pytz                      2025.1\n",
      "pywin32                   308\n",
      "pywinpty                  2.0.15\n",
      "PyYAML                    6.0.2\n",
      "pyzmq                     26.3.0\n",
      "qtconsole                 5.6.1\n",
      "QtPy                      2.4.1\n",
      "rapid-table               1.0.3\n",
      "RapidFuzz                 3.12.2\n",
      "rapidocr-onnxruntime      1.4.4\n",
      "rapidocr-paddle           1.4.5\n",
      "rarfile                   4.2\n",
      "referencing               0.36.2\n",
      "regex                     2024.11.6\n",
      "requests                  2.32.3\n",
      "rfc3339-validator         0.1.4\n",
      "rfc3986-validator         0.1.1\n",
      "robust-downloader         0.0.2\n",
      "rpds-py                   0.23.1\n",
      "s3transfer                0.11.4\n",
      "safetensors               0.5.3\n",
      "scikit-image              0.25.2\n",
      "scikit-learn              1.6.1\n",
      "scipy                     1.15.2\n",
      "seaborn                   0.13.2\n",
      "Send2Trash                1.8.3\n",
      "setuptools                75.8.0\n",
      "shapely                   2.0.7\n",
      "simsimd                   6.2.1\n",
      "sip                       6.10.0\n",
      "six                       1.16.0\n",
      "sniffio                   1.3.0\n",
      "soupsieve                 2.5\n",
      "stack-data                0.6.3\n",
      "stringzilla               3.12.3\n",
      "struct-eqtable            0.3.2\n",
      "sympy                     1.13.1\n",
      "tabulate                  0.9.0\n",
      "tbb                       2021.11.0\n",
      "tensorboard               2.19.0\n",
      "tensorboard-data-server   0.7.2\n",
      "termcolor                 2.5.0\n",
      "terminado                 0.18.1\n",
      "thop                      0.1.1.post2209072238\n",
      "threadpoolctl             3.5.0\n",
      "tifffile                  2025.2.18\n",
      "timm                      0.9.16\n",
      "tinycss2                  1.4.0\n",
      "tokenizers                0.19.1\n",
      "tomli                     2.0.1\n",
      "torch                     2.3.1+cu118\n",
      "torchvision               0.18.1+cu118\n",
      "tornado                   6.4.2\n",
      "tqdm                      4.67.1\n",
      "traitlets                 5.14.3\n",
      "transformers              4.42.4\n",
      "types-python-dateutil     2.9.0.20241206\n",
      "typing_extensions         4.12.2\n",
      "tzdata                    2025.1\n",
      "ultralytics               8.3.86\n",
      "ultralytics-thop          2.0.14\n",
      "unimernet                 0.2.3\n",
      "uri-template              1.3.0\n",
      "urllib3                   2.3.0\n",
      "visualdl                  2.5.3\n",
      "Wand                      0.6.13\n",
      "wcwidth                   0.2.13\n",
      "webcolors                 24.11.1\n",
      "webdataset                0.2.111\n",
      "webencodings              0.5.1\n",
      "websocket-client          1.8.0\n",
      "Werkzeug                  3.1.3\n",
      "wheel                     0.45.1\n",
      "widgetsnbextension        4.0.13\n",
      "win-inet-pton             1.1.0\n",
      "win32_setctime            1.2.0\n",
      "xxhash                    3.5.0\n",
      "yacs                      0.1.8\n",
      "yarl                      1.18.3\n"
     ]
    }
   ],
   "source": [
    "!pip list"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a9a60b6-7b6d-4a19-9171-33324104b4b0",
   "metadata": {},
   "source": [
    "## 安装Mineru"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5c080ab6-1bde-4e01-bf33-de9c848d6d96",
   "metadata": {},
   "source": [
    "## Mineru使用"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "6a05c99a-82bd-42b6-b821-57d93d4bad83",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\u001b[93mimport tensorrt_llm failed, if do not use tensorrt, ignore this message\u001b[0m\n",
      "\u001b[93mimport lmdeploy failed, if do not use lmdeploy, ignore this message\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "import os\n",
    "import logging\n",
    "from magic_pdf.data.data_reader_writer import FileBasedDataWriter, FileBasedDataReader\n",
    "from magic_pdf.data.dataset import PymuDocDataset\n",
    "from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze\n",
    "from magic_pdf.config.enums import SupportedPdfParseMethod\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "2dbe1f49-44e1-4790-a0c6-789d6b113664",
   "metadata": {},
   "outputs": [],
   "source": [
    "logging.basicConfig(\n",
    "    level=logging.INFO,\n",
    "    format='%(asctime)s - %(levelname)s - %(message)s',\n",
    "    datefmt='%Y-%m-%d %H:%M:%S'\n",
    ")\n",
    "logger = logging.getLogger('pdf_processor')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "a9732857-3a36-4660-a257-c7e28019ab8f",
   "metadata": {},
   "outputs": [],
   "source": [
    "def process_pdf(pdf_path, output_dir):\n",
    "    \"\"\"\n",
    "    Process a PDF file and generate various output files.\n",
    "\n",
    "    Args:\n",
    "        pdf_path: Path to the PDF file\n",
    "        output_dir: Directory where outputs will be saved\n",
    "    \"\"\"\n",
    "    pdf_filename = os.path.basename(pdf_path)\n",
    "    base_filename = os.path.splitext(pdf_filename)[0]\n",
    "\n",
    "    logger.info(f\"Processing PDF: {pdf_filename}\")\n",
    "\n",
    "    # Prepare directory structure\n",
    "    images_dir_path = os.path.join(output_dir, \"images\")\n",
    "    images_dir_name = os.path.basename(images_dir_path)\n",
    "\n",
    "    os.makedirs(images_dir_path, exist_ok=True)\n",
    "    logger.debug(f\"Created images directory: {images_dir_path}\")\n",
    "\n",
    "    # Initialize file writers\n",
    "    image_writer = FileBasedDataWriter(images_dir_path)\n",
    "    md_writer = FileBasedDataWriter(output_dir)\n",
    "\n",
    "    # Read PDF content\n",
    "    pdf_reader = FileBasedDataReader(\"\")\n",
    "    pdf_bytes = pdf_reader.read(pdf_path)\n",
    "    logger.debug(f\"Read {len(pdf_bytes)} bytes from {pdf_filename}\")\n",
    "\n",
    "    # Process PDF\n",
    "    dataset = PymuDocDataset(pdf_bytes)\n",
    "    pdf_type = dataset.classify()\n",
    "    logger.info(f\"Detected PDF type: {pdf_type}\")\n",
    "\n",
    "    # Apply appropriate processing based on PDF type\n",
    "    if pdf_type == SupportedPdfParseMethod.OCR:\n",
    "        logger.info(f\"Using OCR mode for {pdf_filename}\")\n",
    "        inference_result = dataset.apply(doc_analyze, ocr=True)\n",
    "        processing_result = inference_result.pipe_ocr_mode(image_writer)\n",
    "    else:\n",
    "        logger.info(f\"Using text mode for {pdf_filename}\")\n",
    "        inference_result = dataset.apply(doc_analyze, ocr=False)\n",
    "        processing_result = inference_result.pipe_txt_mode(image_writer)\n",
    "\n",
    "    # Generate output files\n",
    "    logger.debug(\"Generating output files\")\n",
    "    model_pdf_path = os.path.join(output_dir, \"model.pdf\")\n",
    "    inference_result.draw_model(model_pdf_path)\n",
    "    logger.debug(f\"Created model visualization: {model_pdf_path}\")\n",
    "\n",
    "    model_inference_result = inference_result.get_infer_res()\n",
    "\n",
    "    layout_pdf_path = os.path.join(output_dir, \"layout.pdf\")\n",
    "    processing_result.draw_layout(layout_pdf_path)\n",
    "    logger.debug(f\"Created layout visualization: {layout_pdf_path}\")\n",
    "\n",
    "    spans_pdf_path = os.path.join(output_dir, \"spans.pdf\")\n",
    "    processing_result.draw_span(spans_pdf_path)\n",
    "    logger.debug(f\"Created spans visualization: {spans_pdf_path}\")\n",
    "\n",
    "    # Generate markdown content\n",
    "    markdown_content = processing_result.get_markdown(images_dir_name)\n",
    "    markdown_path = f\"{base_filename}.md\"\n",
    "    processing_result.dump_md(md_writer, markdown_path, images_dir_name)\n",
    "    logger.info(f\"Created markdown file: {markdown_path}\")\n",
    "\n",
    "    # Generate content list\n",
    "    content_list = processing_result.get_content_list(images_dir_name)\n",
    "    processing_result.dump_content_list(md_writer, \"content_list.json\", images_dir_name)\n",
    "    logger.debug(\"Created content list JSON\")\n",
    "\n",
    "    # Generate middle JSON\n",
    "    middle_json = processing_result.get_middle_json()\n",
    "    processing_result.dump_middle_json(md_writer, \"middle.json\")\n",
    "    logger.debug(\"Created middle JSON file\")\n",
    "\n",
    "    logger.info(f\"Successfully processed {pdf_filename}\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "b817a8bc-153a-426f-8a60-38d6d6a018c7",
   "metadata": {},
   "outputs": [],
   "source": [
    "pdf_path=\"../data/docs/2503.04697v1.pdf\"\n",
    "output_dir=\"../output/pdf_parse\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "ebd4db70-9f57-43d9-b498-b76aaa6d5250",
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "\u001b[32m2025-03-18 20:41:51.254\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.data.dataset\u001b[0m:\u001b[36m__init__\u001b[0m:\u001b[36m156\u001b[0m - \u001b[1mlang: None\u001b[0m\n",
      "\u001b[32m2025-03-18 20:41:52.967\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.libs.pdf_check\u001b[0m:\u001b[36mdetect_invalid_chars\u001b[0m:\u001b[36m67\u001b[0m - \u001b[1mcid_count: 0, text_len: 30225, cid_chars_radio: 0.0\u001b[0m\n",
      "\u001b[32m2025-03-18 20:41:52.979\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.doc_analyze_by_custom_model\u001b[0m:\u001b[36mdoc_analyze\u001b[0m:\u001b[36m180\u001b[0m - \u001b[1mgpu_memory: 24 GB, batch_ratio: 8\u001b[0m\n",
      "\u001b[32m2025-03-18 20:41:55.899\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.batch_analyze\u001b[0m:\u001b[36m__call__\u001b[0m:\u001b[36m74\u001b[0m - \u001b[1mlayout time: 2.0, image num: 20\u001b[0m\n",
      "\u001b[32m2025-03-18 20:41:57.888\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.batch_analyze\u001b[0m:\u001b[36m__call__\u001b[0m:\u001b[36m85\u001b[0m - \u001b[1mmfd time: 1.99, image num: 20\u001b[0m\n",
      "\u001b[32m2025-03-18 20:42:05.186\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.batch_analyze\u001b[0m:\u001b[36m__call__\u001b[0m:\u001b[36m100\u001b[0m - \u001b[1mmfr time: 7.3, image num: 139\u001b[0m\n",
      "\u001b[32m2025-03-18 20:42:13.525\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.batch_analyze\u001b[0m:\u001b[36m__call__\u001b[0m:\u001b[36m195\u001b[0m - \u001b[1mdet time: 5.39, image num: 232\u001b[0m\n",
      "\u001b[32m2025-03-18 20:42:13.526\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.batch_analyze\u001b[0m:\u001b[36m__call__\u001b[0m:\u001b[36m197\u001b[0m - \u001b[1mtable time: 2.76, image num: 3\u001b[0m\n",
      "\u001b[32m2025-03-18 20:42:13.856\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.doc_analyze_by_custom_model\u001b[0m:\u001b[36mdoc_analyze\u001b[0m:\u001b[36m235\u001b[0m - \u001b[1mgc time: 0.3\u001b[0m\n",
      "\u001b[32m2025-03-18 20:42:13.858\u001b[0m | \u001b[1mINFO    \u001b[0m | \u001b[36mmagic_pdf.model.doc_analyze_by_custom_model\u001b[0m:\u001b[36mdoc_analyze\u001b[0m:\u001b[36m239\u001b[0m - \u001b[1mdoc analyze time: 20.88, speed: 0.96 pages/second\u001b[0m\n"
     ]
    }
   ],
   "source": [
    "process_pdf(pdf_path,output_dir)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bb00d953-5660-4896-b028-dcdf7db035f0",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
